/*
 * Copyright 2011 Carnegie Mellon University
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *  
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, 
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package edu.cmu.lti.ritesdk.rite4qa;

import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import edu.cmu.lti.ritesdk.AnalyzedDataSet;
import edu.cmu.lti.ritesdk.AnalyzedTextPair;
import edu.cmu.lti.ritesdk.DataSet;
import edu.cmu.lti.ritesdk.EvaluationResult;

/**
 * Evaluates RITE4QA runs using the QA metric (TOP@N accuracy and MRR@N) 
 * 
 * @author Hideki Shima
 *
 */
public class EvaluatorForRITE4QA {

  //For experimental purpose
  // default value: false
  private static boolean random = false; // When true, randomize system labels & conf
  private static boolean systemAllY = false; // When true, insert Y for all system labels with 1 confidence 
  private static boolean highConfNBetter = false; // When true, higher confidence N means closer to Y label
  
  public final static String MRR5 = "MRR5";
  public final static String TOP1 = "TOP1";

  private static NumberFormat nf4 = NumberFormat.getInstance();
  private static NumberFormat nf6 = NumberFormat.getInstance();
  static {
    nf4.setMaximumFractionDigits(4);
    nf4.setMinimumFractionDigits(4);
    nf6.setMaximumFractionDigits(6);
    nf6.setMinimumFractionDigits(6);
  }

  /**
   * The main method of the class
   * @param d analyzed data set
   * @param goldData gold standard data
   * @return evaluation result
   */
  public static EvaluationResult evaluate( AnalyzedDataSet d, DataSet goldData ) {

    Map<String,List<AnalyzedTextPair>> analyzedPairById 
     = getAnalyzedPairById(d,random,systemAllY);
    
    //Per id log
    Map<String,String> log = new LinkedHashMap<String,String>();
    
    //key: CLQA Topic ID, value: map (key: rank, value: judgment list);
    Map<String,Map<Integer,List<Boolean>>> judgments = extractSystemData(analyzedPairById, log);

    Set<String> goldTopicIds = new LinkedHashSet<String>();
    for ( String id : goldData.keySet() ) {
      String clqaId = id.substring(0,id.indexOf("_"));
      goldTopicIds.add(clqaId);
    }
    
    EvaluationResult eval = evaluate(judgments, goldTopicIds, log);
    return eval;
  }
  
  public static Map<String,List<AnalyzedTextPair>> getAnalyzedPairById( 
          AnalyzedDataSet d, boolean randomResultMode, boolean systemAllYMode ) {
    Map<String,List<AnalyzedTextPair>> result = new LinkedHashMap<String,List<AnalyzedTextPair>>();
    for ( String id : d.keySet() ) {
      String clqaId = id.substring(0,id.indexOf("_"));
      List<AnalyzedTextPair> list = result.get(clqaId);
      if (list==null) {
        list = new ArrayList<AnalyzedTextPair>();
      }
      AnalyzedTextPair pair = d.get(id);
      if (randomResultMode) {
        pair.setSystemLabel(Math.random()>0.5?"Y":"N");
        pair.setConfidence(Math.random());// use different rand as above
      } else if (systemAllYMode) {
        pair.setSystemLabel("Y");
        pair.setConfidence(1);// use different rand as above
      }
      list.add( pair );
      result.put(clqaId, list);
    }
    return result;
  }
  
  public static Map<String,Map<Integer,List<Boolean>>> extractSystemData( 
          Map<String,List<AnalyzedTextPair>> analyzedPairById,
          Map<String,String> log ) {
    Map<String,Map<Integer,List<Boolean>>> result = new LinkedHashMap<String,Map<Integer,List<Boolean>>>();
    
    for ( String clqaId : analyzedPairById.keySet() ) {
      int rankPerTopic = 0;
      double prevConf = Double.POSITIVE_INFINITY;
      List<AnalyzedTextPair> ySorted = new ArrayList<AnalyzedTextPair>();
      List<AnalyzedTextPair> nSorted = new ArrayList<AnalyzedTextPair>();
      for ( int i=0; i<analyzedPairById.get(clqaId).size(); i++ ) {
        AnalyzedTextPair pair = analyzedPairById.get(clqaId).get(i); // pair from system
        if (pair.getTextPair()==null) continue;
        String systemLabel = pair.getSystemLabel();
        if (systemLabel.equals("Y")) {
          ySorted.add(pair);
        } else if (systemLabel.equals("N")) {
          nSorted.add(pair);
        }
      }
      Collections.sort(ySorted, new YComparator());
      Collections.sort(nSorted, highConfNBetter?new YComparator():new NComparator());

      Map<Integer,List<Boolean>> rankToJudgments = result.get(clqaId);
      if (rankToJudgments==null) {
        rankToJudgments = new LinkedHashMap<Integer,List<Boolean>>();
      }
      for ( int i=0; i<ySorted.size(); i++ ) {
        AnalyzedTextPair a = ySorted.get(i);
        if (Math.abs(prevConf - a.getConfidence()) > 1.0e-6) {
          int prevJudgmentSize = (rankToJudgments.get(rankPerTopic)!=null) 
          ? rankToJudgments.get(rankPerTopic).size() : 1;
          rankPerTopic += prevJudgmentSize;
        }
        List<Boolean> judgments = rankToJudgments.get(rankPerTopic);
        if ( judgments == null ) {
          judgments = new ArrayList<Boolean>();
        }
        boolean jugment = a.getTextPair().getGoldStandardLabel().equals("Y");
        judgments.add(jugment);
        rankToJudgments.put(rankPerTopic, judgments);
        prevConf = a.getConfidence();
      }
      prevConf = Double.POSITIVE_INFINITY;
      for ( int i=0; i<nSorted.size(); i++ ) {
        AnalyzedTextPair a = nSorted.get(i);
        if (Math.abs(prevConf - a.getConfidence()) > 1.0e-6) {
          int prevJudgmentSize = (rankToJudgments.get(rankPerTopic)!=null) 
          ? rankToJudgments.get(rankPerTopic).size() : 1;
          rankPerTopic += prevJudgmentSize;
        }
        List<Boolean> judgments = rankToJudgments.get(rankPerTopic);
        if ( judgments == null ) {
          judgments = new ArrayList<Boolean>();
        }
        boolean jugment = a.getTextPair().getGoldStandardLabel().equals("Y");
        judgments.add(jugment);
        rankToJudgments.put(rankPerTopic, judgments);
        prevConf = a.getConfidence();
      }
      result.put(clqaId, rankToJudgments);
      StringBuilder sb = new StringBuilder();
      for ( AnalyzedTextPair a : ySorted ) {
        if (sb.length()>0) sb.append(", ");
        sb.append(a.getTextPair().getGoldStandardLabel()+" ");
        sb.append(a.getSystemLabel()+" ");
        sb.append(nf6.format(a.getConfidence()));
      }
      for ( AnalyzedTextPair a : nSorted ) {
        if (sb.length()>0) sb.append(", ");
        sb.append(a.getTextPair().getGoldStandardLabel()+" ");
        sb.append(a.getSystemLabel()+" ");
        sb.append(nf6.format(a.getConfidence()));
      }
      log.put(clqaId, sb.toString());
    }
    return result;
  }
  
  private static EvaluationResult evaluate(
          Map<String,Map<Integer,List<Boolean>>> judgmentData,
          Set<String> goldTopicIds, 
          Map<String,String> log) {
    StringBuilder sbSummary  = new StringBuilder();
    StringBuilder sbDetailed = new StringBuilder();
    
    sbDetailed.append("---- Evaluation log 1 ----\n\n");
    sbDetailed.append("Legend: QA_ID ':' ( true_label system_label system_confidence )* ... sorted by system output\n");
    sbDetailed.append("# Sorting criteria: Y with high conf >= Y with low conf >= N with low conf >= N with high conf\n\n");
    for (String clqaId : goldTopicIds) {
      String prevLog = log.get(clqaId);
      sbDetailed.append(clqaId+": "+prevLog+"\n");
    }
    sbDetailed.append("\n\n---- Evaluation log 2 ----\n\n");
    sbDetailed.append("Legend: QA_ID ':' ( rank '=[' is_true_label_Y ']' )* ... \n");
    sbDetailed.append("# Output with same label and confidence belongs to the same rank.\n\n");
    for (String clqaId : goldTopicIds) {
      Map<Integer,List<Boolean>> rankToJudgments = judgmentData.get(clqaId);
      sbDetailed.append(clqaId+": "+rankToJudgments+"\n");
    }
    sbDetailed.append("\n\n---- Evaluation log 3 ----\n\n");
    sbDetailed.append("Legend: QA_ID ':' highest_rank_of_correct reciprocal_rank top1_credit\n");
    sbDetailed.append("# Rank of highest correct answer is tie-broken, e.g. rank 2 tied between 3 outputs -> (2+3+4)/3 = 3\n");
    sbDetailed.append("# MRR5 is the mean of reciprocal_rank.\n");
    sbDetailed.append("# Top1 credit is also tie-broken, e.g. three outputs tied at rank 1 and two are correct -> 2/3.\n");
    sbDetailed.append("# Top1 is the Top1 credit mean of top1 credit.\n\n");
    
    double cumulativeMrrGain = 0;
    double cmuulativeTop1 = 0;
    
    for (String clqaId : goldTopicIds) {
      double highestCorrectRank = 0D;
      double rr = 0D;
      double top1Credit = 0D;
      
      Map<Integer,List<Boolean>> rankToJudgments = judgmentData.get(clqaId);
      if (rankToJudgments!=null && rankToJudgments.get(1)!=null) {
        int numCorrectAt1 = 0;
        for ( boolean j : rankToJudgments.get(1) ) {
          if (j) numCorrectAt1++;
        }
        top1Credit = (double)numCorrectAt1/(double)rankToJudgments.get(1).size();
        cmuulativeTop1 += top1Credit;
        for ( int rank : rankToJudgments.keySet()) {
          List<Boolean> judgments = rankToJudgments.get(rank);
          boolean atLeastOneTrue = false;
          for ( boolean j : judgments ) {
            atLeastOneTrue = j || atLeastOneTrue;
          }
          if (atLeastOneTrue) {
            double adjustedRank = 0D;
            for ( int i=0; i<judgments.size(); i++ ) {
              adjustedRank += (double)(rank+i);
            }
            adjustedRank /= (double)judgments.size();
            
            highestCorrectRank = adjustedRank;
            rr = 1D/highestCorrectRank;
            cumulativeMrrGain += rr;
            break;
          }
        }
      }
      sbDetailed.append( clqaId+": "+highestCorrectRank+" "+nf4.format(rr)+" "+nf4.format(top1Credit)+"\n" );
    }
    int N = goldTopicIds.size();
    double top1 = cmuulativeTop1 / (double)N;
    double mrr5 = cumulativeMrrGain / (double)N;
    Map<String,String> scores = new LinkedHashMap<String,String>();
    scores.put(TOP1, nf4.format(top1));
    scores.put(MRR5, nf4.format(mrr5));
    String detailed = "---- Score ----\n\n"
      +"Top1 = "+nf4.format(top1)+"\n"
      +"MRR5 = "+nf4.format(mrr5)+"\n\n"
      +sbDetailed.toString();
    EvaluationResult result = new EvaluationResult( scores, sbSummary.toString(), detailed );
    return result;
  }

  public static class YComparator implements Comparator<AnalyzedTextPair> {
    @Override
    public int compare(AnalyzedTextPair o1, AnalyzedTextPair o2) {
      return Double.compare(o2.getConfidence(), o1.getConfidence());
    }
  }
  
  public static class NComparator implements Comparator<AnalyzedTextPair> {
    @Override
    public int compare(AnalyzedTextPair o1, AnalyzedTextPair o2) {
      return Double.compare(o1.getConfidence(), o2.getConfidence());
    }
  }

  public static class Randomizer<T> implements Comparator<T> {
    @Override
    public int compare(T o1, T o2) {
      return Math.random()>0.5D?1:-1;
    }
  }
  
}
